import argparse
import sys
import os
import os.path as osp
import glob
from pipes import quote
from multiprocessing import Pool, current_process
from tqdm import tqdm
import mmcv

FRAME_RATE = 29.97
# >ffmpeg -i UCF-101\ApplyEyeMakeup\v_ApplyEyeMakeup_g01_c01.avi -r 25 img_%04d.png
# >ffmpeg -i videos\ApplyEyeMakeup\v_ApplyEyeMakeup_g01_c01.avi -r 25 img_%04d.png

def dump_frames(vid_item):
    full_path, done_vid_path, num_vid_frame = vid_item
    vid_name = vid_path.split('.')[0]
    out_full_path = osp.join(args.out_dir, vid_name)
    try:
        os.mkdir(out_full_path)
    except OSError:
        pass
    save_path = '{}/img_%04d.png'.format(out_full_path)
    command = ['ffmpeg', 
               '-i', '"{}"'.format(full_path),
               '-r', '{}'.format(FRAME_RATE),
               '"{}"'.format(save_path)]
    command = ' '.join(command)
    status = os.system(command)
    if status is not 0:
        print('Fail to extract frames from {}'.format(vid_name))
    # vr = mmcv.VideoReader(full_path)
    # for i in range(len(vr)):
    #     if vr[i] is not None:
    #         mmcv.imwrite(
    #             vr[i], '{}/img_{:05d}.jpg'.format(out_full_path, i + 1))
    #     else:
    #         print('[Warning] length inconsistent!'
    #               'Early stop with {} out of {} frames'.format(i + 1, len(vr)))
    #         break
    # print('{} done with {} frames'.format(vid_name, len(vr)))
    sys.stdout.flush()
    return True

def check_frames(vid_item):
    full_path, done_vid_path, num_img_frame = vid_item
    vid_name = done_vid_path.split('/')[-1]
    check_command_prefix = "ffprobe -v error -count_frames -select_streams v:0 -show_entries stream=nb_read_frames -of default=nokey=1:noprint_wrappers=1"
    check_command = ' '.join([check_command_prefix, full_path])
    f = os.popen(check_command)
    num_vid_frame = int(f.readlines()[0]) + 1
    f.close()
    check_command_prefix = "ffprobe -v error -select_streams v -of default=noprint_wrappers=1:nokey=1 -show_entries stream=r_frame_rate"
    check_command = ' '.join([check_command_prefix, full_path])
    f = os.popen(check_command)
    sys_out = f.readlines()[0]
    fps = float(sys_out.split('/')[0]) / float(sys_out.split('/')[1])
    f.close()
    error = num_vid_frame - num_img_frame
    if abs(error) > 1:
        print("{} has {} frames, but there are {} imgs instead. Error is {}. FPS-{:.2f}".format(vid_name, num_vid_frame, 
                                                                                                num_img_frame, error, fps))
        # os.system("rm {}".format(done_vid_path+'/*.png'))
        # save_path = '{}/img_%04d.png'.format(done_vid_path)
        # command = ['ffmpeg', 
        #        '-i', '"{}"'.format(full_path),
        #        '-r', '{}'.format(fps),
        #        '"{}"'.format(save_path)]
        # command = ' '.join(command)
        # status = os.system(command)
        # len(glob.glob(p + '/*.png')
    # save_path = '{}/img_%04d.png'.format(done_vid_path)
    # command = ['ffmpeg', 
    #            '-i', '"{}"'.format(full_path),
    #            '-r', '{}'.format(fps),
    #            '"{}"'.format(save_path)]
    # command = ' '.join(command)
    # status = os.system(command)

    # vr = mmcv.VideoReader(full_path)
    # for i in range(len(vr)):
    #     if vr[i] is not None:
    #         mmcv.imwrite(
    #             vr[i], '{}/img_{:05d}.jpg'.format(out_full_path, i + 1))
    #     else:
    #         print('[Warning] length inconsistent!'
    #               'Early stop with {} out of {} frames'.format(i + 1, len(vr)))
    #         break
    # print('{} done with {} frames'.format(vid_name, len(vr)))
    sys.stdout.flush()
    return True


def parse_args():
    parser = argparse.ArgumentParser(description='extract optical flows')
    parser.add_argument('src_dir', type=str)
    parser.add_argument('out_dir', type=str)
    parser.add_argument('--level', type=int,
                        choices=[1, 2],
                        default=2)
    parser.add_argument('--num_worker', type=int, default=16)
    parser.add_argument('--flow_type', type=str,
                        default=None, choices=[None, 'tvl1', 'warp_tvl1'])
    parser.add_argument('--df_path', type=str,
                        default='../../mmaction/third_party/dense_flow')
    parser.add_argument("--out_format", type=str, default='dir',
                        choices=['dir', 'zip'], help='output format')
    parser.add_argument("--ext", type=str, default='avi',
                        choices=['avi', 'mp4'], help='video file extensions')
    parser.add_argument("--new_width", type=int, default=0,
                        help='resize image width')
    parser.add_argument("--new_height", type=int,
                        default=0, help='resize image height')
    parser.add_argument("--num_gpu", type=int, default=8, help='number of GPU')
    parser.add_argument("--resume", action='store_true', default=False,
                        help='resume optical flow extraction '
                        'instead of overwriting')
    args = parser.parse_args()

    return args


if __name__ == '__main__':
    args = parse_args()

    print('Reading videos from folder: ', args.src_dir)
    print('Extension of videos: ', args.ext)
    if args.level == 2:
        fullpath_list = sorted(glob.glob(args.src_dir + '/*/*.' + args.ext))
        done_fullpath_list = sorted(glob.glob(args.out_dir + '/*/*'))
    elif args.level == 1:
        fullpath_list = glob.glob(args.src_dir + '/*.' + args.ext)
        done_fullpath_list = glob.glob(args.out_dir + '/*')
    print('Total number of videos found: ', len(fullpath_list))
    assert len(fullpath_list) == len(done_fullpath_list), "Video dirs are not enough"
    # num_video_frame_list = list(map(lambda p: len(glob.glob(p + '/*.png')), done_fullpath_list))
    num_video_frame_list = list(map(lambda p: len(glob.glob(p + '/*.jpg')), done_fullpath_list))
    # print(fullpath_list[-10:])
    print(len(done_fullpath_list))
    # print(num_video_frame_list[-10:])
    if args.level == 2:
        vid_list = list(map(lambda p: osp.join(
            '/'.join(p.split('/')[-2:])), fullpath_list))
        # print(vid_list)
    elif args.level == 1:
        vid_list = list(map(lambda p: p.split('/')[-1], fullpath_list))
    
    # ffprobe -v error -count_frames -select_streams v:0 -show_entries stream=nb_read_frames -of default=nokey=1:noprint_wrappers=1 videos/ApplyEyeMakeup/v_ApplyEyeMakeup_g01_c03.avi

    pool = Pool(args.num_worker)

    pool.map(check_frames, zip(
        fullpath_list, done_fullpath_list, num_video_frame_list))
    # for item in tqdm(zip(fullpath_list, done_fullpath_list, num_video_frame_list)):
    #     _ = check_frames(item)